# from backtrader import indicator
import pandas as pd
from pymongo import MongoClient
from datetime import datetime
import time
import math
import backtrader as bt
import backtrader.indicators as btind
import backtrader.feeds as btfeeds
import talib
import numpy as np
from target.indicator_ma import *
objJson = {
    "_id": 'uuidstr',
    "adjustflag": 'adjustflag',
    "code": 'code',
    'start': 'start',
    'end': 'end',
    'startcash': 'startcash',
    'endtcash': 'cerebro.broker.getvalue()',
    'com': 'com',
    'qts': 'qts',
    'basicsData': [],  # 基础数据
    'isbuyArr': [],  # 买入一笔 买入价格
    'issellArr': [],  # 卖出一笔  卖出价格
    'pnlcommArr': [],  # 一笔交易后的盈亏
    'currentCapital' : [],  #资金走势
    'indicatorLines':{}, #用户所选指标线

    'annualReturn':[],  #年度回报
    'Calmar':[],  #卡码率
    'DrawDown':{},  #回撤
    'sharperatio':math.nan,  #夏普比率
    'sqnIndex':{},  #SQN指数
    'TimeReturn':[],  #基于时间范围内回报
}


# 创建策略继承bt.strategy
class TestStrategy(bt.Strategy):
    params = (
        ('maperiod', 15),
        # 判断是否输出该日志
        ('printlog', False),
    )

    #用户代码init
    def init(self):
        # 需要设置计算MACD的相关参数参数
        self.short = 12
        self.long = 26
        self.smoothperiod = 9
        self.observation = 100
        
    
    #用户策略代码
    def on_bar(self):
        if self.order:
            return
        self.cash = self.get_cash()   #获取现金
        self.dataclose = self.datas[0].close  #获取价格
        self.dates = self.datas[0].datetime.date(0).isoformat()  #当前日

        self.priceArr = self.optainData(data = self.datas[0].close,ago = 0,size = self.observation)
        if len(self.priceArr) < 100:
            return
        else: 

            macd, signal, hist = talib.MACD(self.priceArr, self.short,
                                    self.long, self.smoothperiod)

            # 策略逻辑
            # 卖出逻辑 macd下穿signal
            if macd[-1] - signal[-1] < 0 and macd[-2] - signal[-2] > 0:
                # 有持仓进行清仓
                if self.position:
                    self.open_close()
            # 买入逻辑  macd上穿signal
            if macd[-1] - signal[-1] > 0 and macd[-2] - signal[-2] < 0:
                # 空仓买入股票
                if not self.position:
                   self.open_buy(size=100)

    #初始化部分指标
    def __init__(self):

        #调用用户初始化代码
        self.init()
        
        #指标进行注册
        self.indicators = {}
        # 跟踪挂单
        self.order = None
        # 买入价格和手续费
        self.buyprice = None
        self.buycomm = None
        
         # 初始化注册指标
        for param in indicator_params:
            if param['code'] == 'ma':
                indicator = bt.indicators.MovingAverageSimple(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'ema':
                indicator = bt.indicators.ExponentialMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'smma':
                indicator = bt.indicators.SmoothedMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'wma':
                indicator = bt.indicators.WeightedMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'dema':
                indicator = bt.indicators.DoubleExponentialMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'kama':
                indicator = bt.indicators.AdaptiveMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'zlema':
                indicator = bt.indicators.ZeroLagExponentialMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'hma':
                indicator = bt.indicators.HullMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'zlind':
                indicator = bt.indicators.ZeroLagIndicator(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

            elif param['code'] == 'dma':
                indicator = bt.indicators.DicksonMovingAverage(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator
            
            elif param['code'] == 'highest':
                indicator = bt.indicators.Highest(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator
            
            elif param['code'] == 'lowest':
                indicator = bt.indicators.Lowest(
                    param['data'], period=param['params'][0])
                self.indicators[param['name']] = indicator
                # 修改注册表
                value = indicator_registry[param['name']]
                value['indicator'] = indicator

        # 初始化指标lines对象(用于存储绘图数据)
        for item in indicator_user_arr:
            objJson['indicatorLines'][item] = []

    # 策略核心代码
    def next(self):

        current_date = self.datas[0].datetime.date(0).isoformat()  #当前日期
        current_capital = self.broker.getvalue()  #当前剩余资金
        # print(self.position.size)
        objJson['currentCapital'].append({'date':current_date,'value':round(current_capital, 2)})  #保存每一天的资金走势

        #indicator_user_arr：当前用户所选指标列表
        for item in indicator_user_arr:
            #每一个指标线的值进行保存
            objJson['indicatorLines'][item].append({'date':current_date,'value':round(self.indicators[item][0],2)})
        
        # print('日期:{}---ma5:{}---ma50:{}'.format(current_date,self.ma5[0],self.ma50[0]))
        
        #调用用户策略
        self.on_bar()

        
    
    #买入
    def open_buy(self,size=None):
        self.order = self.buy(size=size)
    #卖出
    def open_sell(self,size=None):
        self.order = self.sell(size=size)
    #平仓
    def open_close(self):
        self.order = self.close()
    #取消订单
    def open_cancel(self,size=None):
        self.order = self.cancel(size=size)
    #获取资金
    def get_cash(self):
        return self.broker.getvalue()
    #获取某个时间段内的数据
    # data  获取指定列的数据
    # ago 0：从当前时间点往前取数据（默认为0）
        # 1：从当前时间点的前一天开始
    #size  时间点数据（默认为10）
    def optainData(self,data,ago=0,size = 10):
        return np.array(data.get(ago=ago, size=size))



    # 日志输出
    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            # 记录策略的执行日志
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')

    # 订单状态通知，买入卖出都是下单
    def notify_order(self, order):
        # 提交了/接受了,  买/卖订单什么都不做
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 检查一个订单是否完成
        # 注意:当资金不足时，broker会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                print('日期:{}---买入价格:{}---数量:{}---买入手续费{}'.format(self.dates,order.executed.price,order.executed.size,order.executed.comm))
                objJson['isbuyArr'].append(
                    {'date': self.datas[0].datetime.date(0).isoformat(), 'value': round(order.executed.price, 4),'size':order.executed.size,'comm':round(order.executed.comm,4)})
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            elif order.issell():
                print('日期:{}---卖出价格:{}---数量:{}---卖出手续费{}'.format(self.dates,order.executed.price,order.executed.size,order.executed.comm))
                objJson['issellArr'].append(
                    {'date': self.datas[0].datetime.date(0).isoformat(), 'value': round(order.executed.price, 4),'size':order.executed.size,'comm':round(order.executed.comm,4)})
            self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('订单取消/保证金不足/拒绝', doprint=False)
        # 其他状态记录为：无法挂单
        self.order = None

    # 交易状态通知，一买一卖算交易（交易净利润）
    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        objJson['pnlcommArr'].append(
            {'date': self.datas[0].datetime.date(0).isoformat(), 'value': round(trade.pnlcomm, 2)})

    # 策略结束时，多用于参数调优
    def stop(self):
        # 回测结束存储数据
        # print(objJson)
        self.log('(MA均线： %2d日) 期末总资金 %.2f' %
                 (self.params.maperiod, self.broker.getvalue()), doprint=False)


def main(execution_id, adjustflag, code, start, end='', startcash=1000000, com=0.0015, qts=100):

    client = MongoClient('47.100.19.231', 27018)
    db = client.admin
    db.authenticate("quant", "Qweasd123")
    db = client.quant
    mycol = db.t_strategy_execution_result
    mydict = {'_id': execution_id}
    s = mycol.find(mydict).count()
    if s > 0:
        return execution_id

    # 创建主控制器
    cerebro = bt.Cerebro()
    # 导入策略参数寻优
    cerebro.addstrategy(TestStrategy)
    # cerebro.optstrategy(TestStrategy,maperiod=range(10, 15))
    # 获取数据
    df = selectSotck(adjustflag, code, start, end)
    df.index = pd.to_datetime(df.date)
    df = df[['open', 'high', 'low', 'close', 'volume']]
    # 将数据加载至回测系统
    data = bt.feeds.PandasData(dataname=df)

    cerebro.adddata(data)
    # broker设置资金、手续费
    cerebro.broker.setcash(startcash)
    cerebro.broker.setcommission(commission=com)
    # cerebro.broker.setcommission(PerOrder(buy_cost=0.0003, sell_cost=0.0013, min_cost=5))
    # 设置买入设置，策略，数量
    cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
    # 以发出信号当日收盘价成交
    # cerebro.broker.set_coc(True)
    print('期初总资金: %.2f' % cerebro.broker.getvalue())
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='AnnualReturn')  # 年度回报
    cerebro.addanalyzer(bt.analyzers.Calmar, _name='Calmar')  # 卡玛率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='DW')  # 回撤
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.SQN, _name='SQN')  # SQN指数
    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='TimeReturn')  # 基于时间范围的回报
    result = cerebro.run(maxcpus=1)
    strat = result[0]

    #年度回报
    for item in (strat.analyzers.AnnualReturn.get_analysis()).items():
         objJson['annualReturn'].append(
            {'date': item[0], 'value': round(item[1],4)})
    #卡码率
    for item in (strat.analyzers.Calmar.get_analysis()).items():
        objJson['Calmar'].append(
            {'date':item[0].strftime('%Y-%m-%d'), 'value': round(item[1],4)})
    #回撤
    dw = strat.analyzers.DW.get_analysis()
    objJson['DrawDown']['DrawDownlen'] = dw['len']
    objJson['DrawDown']['DrawDown'] = dw['drawdown']
    objJson['DrawDown']['DrawDownMoney'] = dw['moneydown']
    objJson['DrawDown']['maxDrawDownlen'] = dw['max']['len']
    objJson['DrawDown']['maxDrawDown'] = dw['max']['drawdown']
    objJson['DrawDown']['maxDrawDownMoney'] = dw['max']['moneydown']
    #夏普比率
    SharpeRatio = strat.analyzers.SharpeRatio.get_analysis()
    objJson['sharperatio'] = SharpeRatio['sharperatio']
    #SQN指数
    sqn = strat.analyzers.SQN.get_analysis()
    objJson['sqnIndex']['sqn'] = sqn['sqn']
    objJson['sqnIndex']['trades'] = sqn['trades']
    #基于时间范围内的回报
    for item in (strat.analyzers.TimeReturn.get_analysis()).items():
        objJson['TimeReturn'].append(
            {'date':item[0].strftime('%Y-%m-%d'), 'value': round(item[1],4)})

    print('期末总资金: %.2f' % cerebro.broker.getvalue())
    cerebro.plot()



    objJson['_id'] = execution_id
    objJson['adjustflag'] = adjustflag
    objJson['code'] = code
    objJson['start'] = start
    objJson['end'] = end
    objJson['startcash'] = startcash
    objJson['endtcash'] = cerebro.broker.getvalue()
    objJson['com'] = com
    objJson['qts'] = qts
    # mycol.insert_one(objJson)

    return execution_id


# 时间处理
def loadDate(str):
    date = datetime.strptime(str, '%Y%m%d').strftime('%Y-%m-%d')
    return (date)


# 数据查询
def selectSotck(adjustflag,code, startDate, endDate):
    client = MongoClient('47.100.19.231', 27018)
    db = client.admin
    db.authenticate("quant", "Qweasd123")
    db = client.quant
    mycol = db.testKline

    myquery = {
        "adjustflag": adjustflag,  # 复权类型
        "code": code,  # 股票代码
        "date": {'$gt': startDate, '$lt': endDate}  # 时间区间
    }

    mydoc = mycol.find(myquery)
    data_list = []
    for x in mydoc:
        data_list.append(x)
    objJson['basicsData'] = data_list
    fields = ['date', 'open', 'close', 'high', 'low', 'volume', 'code']
    df = pd.DataFrame(data_list, columns=fields)
    df['date'] = df['date'].map(loadDate)
    test_list = ['open', 'close', 'high', 'low', 'volume']
    df[test_list] = df[test_list].astype('float64')
    return df

main('1','3','sz.000651', '20100916', '20171108')